# model.py - QSVM 模型训练与识别逻辑

import numpy as np
from qiskit_aer import Aer
from qiskit_machine_learning.kernels import QuantumKernel
from qiskit_machine_learning.algorithms import QSVC
from qiskit.circuit.library import ZZFeatureMap
from qiskit.utils import algorithm_globals
from sklearn.preprocessing import normalize
from sklearn.decomposition import PCA
from PIL import Image

# --- QSVM Setup and Training ---

# 设置随机种子
algorithm_globals.random_seed = 42


def train_model():
    """训练用于字符识别的 QSVM 模型"""
    # 设计 8x8 像素的训练数据 (仅 a, b, c)
    train_data = np.zeros((3, 8 * 8))

    # a
    a = np.array([
        [0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 1, 1, 1, 0, 0, 0],
        [0, 1, 0, 0, 0, 1, 0, 0],
        [0, 1, 0, 0, 0, 1, 0, 0],
        [0, 0, 1, 1, 1, 1, 0, 0],
        [0, 0, 0, 0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0]
    ])
    # 注意：这里的训练数据是黑底白字 (0=背景, 1=墨迹)。
    # 如果输入图像反转后是这样，就不需要再反转输入数据。
    # 如果输入图像反转后是白底黑字(0=墨迹, 255=背景)，训练数据或输入数据需要统一。
    # 目前假设QSVC能处理0-1范围的输入，normalize会处理。
    train_data[0] = a.flatten()

    # b
    b = np.array([
        [0, 0, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0, 0, 0],
        [0, 1, 1, 1, 1, 0, 0, 0],
        [0, 1, 0, 0, 0, 1, 0, 0],
        [0, 1, 0, 0, 0, 1, 0, 0],
        [0, 1, 1, 1, 1, 0, 0, 0]
    ])
    train_data[1] = b.flatten()

    # c
    c = np.array([
        [0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 1, 1, 1, 0, 0, 0],
        [0, 1, 0, 0, 0, 1, 0, 0],
        [0, 1, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 1, 0, 0],
        [0, 0, 1, 1, 1, 0, 0, 0]
    ])
    train_data[2] = c.flatten()

    # 使用 PCA 进行降维到 3 维
    pca = PCA(n_components=3)
    # 注意：这里传入的是 0/1 数据，PCA 应该能处理
    train_data_reduced = pca.fit_transform(train_data)
    # 归一化降维后的数据
    train_data_reduced = normalize(train_data_reduced)
    train_labels = [0, 1, 2]  # 0="a", 1="b", 2="c"

    # 设置量子核和 QSVC
    feature_map = ZZFeatureMap(feature_dimension=3, reps=1)
    quantum_instance = Aer.get_backend("aer_simulator")
    quantum_kernel = QuantumKernel(feature_map=feature_map, quantum_instance=quantum_instance)
    qsvc = QSVC(quantum_kernel=quantum_kernel)

    print("开始 QSVM 训练...")
    qsvc.fit(train_data_reduced, train_labels)

    print("🎉 QSVM 模型训练完成！")
    print("支持识别的字母：a, b, c")
    print("您现在可以开始使用手写识别功能了！")
    print("-" * 50)

    return qsvc, pca


def split_image(image):
    """
    动态地将手写图像分割成字符片段。
    包含调试打印、可调参数和针对白色背景的像素反转。
    """
    # --- 可调参数 ---
    # !!! 重要：进行像素反转后，需要根据新的列像素和范围调整阈值 !!!
    # 墨迹现在是 255，背景是 0。
    # 一个合理的初始猜测可能在几百到几千之间，取决于线条粗细和图像高度。
    # 先运行一次，观察反转后的 "Max/Min column sum"，然后设置一个介于 Min (空白列) 和 Max (墨迹列) 之间的值。
    threshold = 1000  # <<<--- 初始猜测值，需要根据输出调整！
    padding = 5       # 在检测到的字符框周围添加的像素数。
    min_char_width = 5 # 被视为字符的最小像素宽度
    min_char_height = 5 # 被视为字符的最小像素高度
    # --- END 可调参数 ---

    print("\n--- 进入 split_image ---")
    width, height = image.size
    print(f"输入图像尺寸: {width}x{height}")

    # 转换为灰度 numpy 数组 (初始时 0=黑, 255=白)
    pixels = np.array(image.convert("L"))
    print(f"像素数组形状: {pixels.shape}")
    print(f"像素值 (反转前样本): Max={np.max(pixels)}, Min={np.min(pixels)}")

    # --- 解决方案：反转像素值 ---
    # 反转图像，使背景为 0，墨迹为 255
    pixels = 255 - pixels
    print(f"像素值 (反转后样本): Max={np.max(pixels)}, Min={np.min(pixels)}")
    # --- END 解决方案 ---


    # 计算列像素和 (现在高值代表墨迹)
    column_sums = np.sum(pixels, axis=0)
    # print(f"列像素和 (前 50): {column_sums[:50]}") # 打印前 50 个和以保持简洁
    print(f"列像素和 (反转后): Max={np.max(column_sums)}, Min={np.min(column_sums)}") # <<<--- 观察这些新值!
    print(f"使用的阈值: {threshold}") # <<<--- 将阈值与新的 Max/Min 和进行比较

    boundaries = []
    in_char = False
    start_char = 0

    for i, value in enumerate(column_sums):
        is_above_threshold = value > threshold
        if is_above_threshold and not in_char:
            # 检测到潜在字符的开始
            start_char = max(0, i - padding)
            in_char = True
            # print(f"  [Debug] 在列 {i} 附近检测到字符开始 (sum={value}). 边界起始设为 {start_char}")
        elif not is_above_threshold and in_char:
            # 检测到字符的结束 (列像素和低于阈值 = 字符间空格)
            end_char = min(width, i + padding)
            detected_width = end_char - start_char
            # print(f"  [Debug] 在列 {i} 附近检测到字符结束 (sum={value}). 边界结束设为 {end_char}. 宽度={detected_width}")
            # 仅当片段足够宽时才添加
            if detected_width > min_char_width: # 简单检查：仅检查是否大于最小宽度
                 boundaries.append((start_char, end_char))
                 # print(f"  [Debug] 添加边界: ({start_char}, {end_char})")
            # else:
            #      print(f"  [Debug] 拒绝边界: 太窄 ({detected_width} <= {min_char_width})")
            in_char = False # 重置标志

    # 处理图像在字符中间结束的情况
    if in_char:
        end_char = min(width, len(column_sums) + padding)
        detected_width = end_char - start_char
        print(f"  [Debug] 图像在字符内结束. 边界结束设为 {end_char}. 宽度={detected_width}")
        if detected_width > min_char_width:
             boundaries.append((start_char, end_char))
             print(f"  [Debug] 添加最终边界: ({start_char}, {end_char})")
        # else:
        #      print(f"  [Debug] 拒绝最终边界: 太窄 ({detected_width} <= {min_char_width})")

    chars = []
    print(f"检测到的水平边界: {boundaries}") # <<<--- 现在仔细检查这个输出！
    print("处理检测到的边界...")
    for i, (left, right) in enumerate(boundaries):
        print(f" 处理边界 {i+1}: 列 {left}-{right} (宽度={right-left})")
        if right - left > min_char_width:
            # 为此片段找到紧密的垂直边界
            # 也使用反转后的像素进行行求和
            char_segment = pixels[:, left:right]
            row_sums = np.sum(char_segment, axis=1)
            # 垂直方向使用相同的阈值？或者如果线条很细，可以使用小一点的？
            # 这里用列阈值的一个分数作为示例，可能需要调整
            vertical_threshold = threshold / 5
            non_empty_rows = np.where(row_sums > vertical_threshold)[0]
            # print(f"   垂直行和 (max={np.max(row_sums)}): {row_sums}")
            # print(f"   非空行索引 (使用阈值 {vertical_threshold}): {non_empty_rows}")

            if len(non_empty_rows) > 0:
                top = max(0, non_empty_rows[0] - padding)
                # numpy切片不包括结束索引，PIL crop对底部/右侧也是排他的。
                # 因此，使用 non_empty_rows[-1] + 1 获取包含墨迹的最后一行，然后添加填充。
                bottom = min(height, non_empty_rows[-1] + 1 + padding)
                detected_height = bottom - top
                print(f"   找到垂直边界: 行 {top}-{bottom} (高度={detected_height})")

                if detected_height > min_char_height:
                    # 使用紧密边界裁剪字符 (注意：需要从原始未反转的图像裁剪，还是反转后的？)
                    # 为了与训练数据匹配（假设是黑底白字），我们可能需要从反转后的像素创建图像
                    # 或者，更简单的是，裁剪原始图像，然后在recognize_word中处理像素值

                    # 方案 A: 裁剪原始图像 (白底黑字)
                    # char_img_pil = image.crop((left, top, right, bottom))

                    # 方案 B: 从反转后的像素创建图像 (黑底白字) - 可能更符合训练数据
                    char_segment_inverted = pixels[top:bottom, left:right] # 从反转的像素中提取
                    char_img_pil = Image.fromarray(char_segment_inverted, mode='L') # 创建PIL图像

                    print(f"   裁剪出的字符图像尺寸: {char_img_pil.size}")

                    # 调整大小为 8x8 以供模型使用
                    try:
                        char_img_resized = char_img_pil.resize((8, 8), Image.Resampling.LANCZOS)
                        chars.append(char_img_resized)
                        print(f"   已调整大小为 8x8 并添加到字符列表。")
                        # 可选：保存中间图像用于调试
                        # char_img_resized.save(f"debug_char_{i+1}.png")
                    except Exception as e:
                        print(f"   错误：调整边界 {i+1} 的图像大小时出错: {e}")

                else:
                     print(f"   跳过片段：太矮 (高度 {detected_height} <= {min_char_height})")
            else:
                 print(f"   跳过片段：垂直方向未找到墨迹 (所有 row_sums <= {vertical_threshold})")
        else:
            print(f" 跳过片段：太窄 (宽度 {right-left} <= {min_char_width})")

    print(f"--- 退出 split_image: 分割成 {len(chars)} 个字符。 ---")
    return chars


def recognize_word(image, qsvc, pca):
    """通过分割、处理和预测每个字符来识别单词。"""
    # 创建副本以避免修改可能需要的全局图像
    img_copy = image.copy()
    char_images = split_image(img_copy) # split_image 现在返回的是处理过的、可能是反转后的字符图像列表

    if not char_images:
        print("分割后未检测到字符。")
        return ""

    predicted_word = ""
    print("逐字符预测:")
    for i, char_img in enumerate(char_images):
        # char_img 现在是分割并调整大小后的 8x8 PIL 图像
        # 它可能是黑底白字 (来自方案B) 或 白底黑字 (来自方案A)
        # 需要确保送入模型的数据格式与训练时一致

        # 确保图像是 'L' 模式 (灰度)
        if char_img.mode != 'L':
            char_img = char_img.convert('L')

        # 将 8x8 PIL 图像转换为扁平化的 numpy 数组
        pixel_data = np.array(char_img).flatten()

        # --- 数据格式统一 ---
        # 检查像素值范围，判断是黑底白字还是白底黑字
        # 假设训练数据是 0=黑, 1=白
        if np.max(pixel_data) > 1: # 如果最大值大于1 (例如 255)，说明是 0-255 范围
            # 将其归一化到 0-1 范围
            # 如果是白底黑字(0=墨, 255=白)，需要反转并归一化 -> (255-pixel)/255
            # 如果是黑底白字(255=墨, 0=白) (来自方案B)，只需归一化 -> pixel/255
            # 我们采用方案B，所以只需归一化
            pixel_data = pixel_data / 255.0
            print(f"  [Debug] 字符 {i+1}: 像素已从 0-255 归一化到 0-1。")
        # 如果已经是 0-1 范围，则无需操作 (理论上不应该发生，因为split_image输出PIL图像)

        # 应用训练数据上拟合的 *相同* PCA 变换
        # PCA 要求输入是 [n_samples, n_features]
        pixel_data_reduced = pca.transform([pixel_data])[0]

        # 对降维后的数据进行归一化 (对于基于距离的核函数很重要)
        pixel_data_reduced = normalize([pixel_data_reduced])[0]

        print(f"字符 {i + 1} 降维后像素数据: {pixel_data_reduced}")

        # 使用训练好的 QSVC 模型进行预测
        try:
            # QSVC 也要求输入是 [n_samples, n_features]
            prediction = qsvc.predict([pixel_data_reduced])[0]
            predicted_char = chr(97 + prediction)  # 0->'a', 1->'b', 2->'c'
            predicted_word += predicted_char
            print(f"字符 {i + 1} 预测为: {predicted_char}")
        except Exception as e:
            print(f"错误：预测字符 {i + 1} 时出错: {e}")
            predicted_word += "?"  # 添加占位符表示错误

    print(f"整体预测结果: {predicted_word}")
    return predicted_word